package org.dromara.common.websocket.interceptor

import cn.dev33.satoken.exception.NotLoginException
import cn.dev33.satoken.stp.StpUtil
import io.github.oshai.kotlinlogging.KotlinLogging
import org.dromara.common.core.domain.model.LoginUser
import org.dromara.common.core.utils.ServletUtils.getParameter
import org.dromara.common.core.utils.ServletUtils.getRequest
import org.dromara.common.core.utils.StringUtils
import org.dromara.common.satoken.utils.LoginHelper
import org.dromara.common.satoken.utils.LoginHelper.getLoginUser
import org.dromara.common.websocket.constant.WebSocketConstants.LOGIN_USER_KEY
import org.springframework.http.server.ServerHttpRequest
import org.springframework.http.server.ServerHttpResponse
import org.springframework.web.socket.WebSocketHandler
import org.springframework.web.socket.server.HandshakeInterceptor

/**
 * WebSocket握手请求的拦截器
 *
 * @author LikeYouDo
 * @date 2025/09/22
 */
class PlusWebSocketInterceptor : HandshakeInterceptor {
    companion object {
        private val log = KotlinLogging.logger { }
    }

    /**
     * WebSocket握手之前执行的前置处理方法
     *
     * @param request    WebSocket握手请求
     * @param response   WebSocket握手响应
     * @param wsHandler  WebSocket处理程序
     * @param attributes 与WebSocket会话关联的属性
     * @return 如果允许握手继续进行，则返回true；否则返回false
     */
    override fun beforeHandshake(
        request: ServerHttpRequest,
        response: ServerHttpResponse,
        wsHandler: WebSocketHandler,
        attributes: MutableMap<String?, Any?>
    ): Boolean {
        try {
            // 检查是否登录 是否有token
            val loginUser = getLoginUser<LoginUser>()

            // 解决 ws 不走 mvc 拦截器问题(cloud 版本不受影响)
            // 检查 header 与 param 里的 clientid 与 token 里的是否一致
            val headerCid = getRequest().getHeader(LoginHelper.CLIENT_KEY)
            val paramCid = getParameter(LoginHelper.CLIENT_KEY)
            val clientId: String = StpUtil.getExtra(LoginHelper.CLIENT_KEY).toString()
            if (!StringUtils.equalsAny(clientId, headerCid, paramCid)) {
                // token 无效
                throw NotLoginException.newInstance(
                    StpUtil.getLoginType(),
                    "-100", "客户端ID与Token不匹配",
                    StpUtil.getTokenValue()
                )
            }

            attributes[LOGIN_USER_KEY] = loginUser
            return true
        } catch (e: NotLoginException) {
            log.error { "WebSocket 认证失败'${e.message}',无法访问系统资源" }
            return false
        }
    }

    /**
     * WebSocket握手成功后执行的后置处理方法
     *
     * @param request   WebSocket握手请求
     * @param response  WebSocket握手响应
     * @param wsHandler WebSocket处理程序
     * @param exception 握手过程中可能出现的异常
     */
    override fun afterHandshake(
        request: ServerHttpRequest,
        response: ServerHttpResponse,
        wsHandler: WebSocketHandler,
        exception: Exception?
    ) {
        // 在这个方法中可以执行一些握手成功后的后续处理逻辑，比如记录日志或者其他操作
    }
}
